import os
import uuid
import wave
import logging
from datetime import datetime

import numpy as np
import torch
from funasr import AutoModel

logger = logging.getLogger(__name__)


class FSMNVAD:
    def __init__(self, config):
        # 如果没有提供配置，使用默认配置
        if config is None:
            config = {}

        print("Initializing FSMNVAD with config:", config)

        # 加载 iic/speech_fsmn_vad_zh-cn-16k-common-pytorch 模型
        self.model = AutoModel(
            model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
            trust_remote_code=True
        )

        # 设置配置参数，使用默认值作为后备
        self.sampling_rate = config.get("sampling_rate", 16000)
        self.threshold = config.get("threshold", 0.5)
        self.frame_length = config.get("frame_length", 320)  # 20ms @ 16kHz
        self.min_silence_duration_ms = config.get("min_silence_duration_ms", 200)

        # 初始化缓存和状态
        self.cache = {}
        self.is_speaking = False
        self.silence_counter = 0
        self.frame_time_ms = (self.frame_length / self.sampling_rate) * 1000

        logger.debug(f"FSMN VAD model initialized with sampling_rate={self.sampling_rate}")

    @staticmethod
    def int2float(sound):
        """
        将 int16 音频数据转换为 float32 格式
        """
        sound = sound.astype(np.float32) / 32768.0
        return sound

    def is_vad(self, data):
        """
        检测音频数据中是否包含语音
        Args:
            data: 音频数据（二进制格式）
        Returns:
            dict: 包含 "start" 或 "end" 键的字典，表示语音开始或结束；如果没有检测到语音，返回 None
        """
        try:
            # 将二进制数据转换为 numpy 数组
            audio_int16 = np.frombuffer(data, dtype=np.int16)
            # 转换为 float32 格式
            audio_float32 = self.int2float(audio_int16)

            # 使用 FSMN VAD 模型进行处理
            vad_result = self.model.generate(
                input=audio_float32,
                cache=self.cache,
                sampling_rate=self.sampling_rate
            )

            # 处理结果并适配系统期望的格式
            if vad_result and isinstance(vad_result, list) and len(vad_result) > 0:
                result = vad_result[0]

                # 检查是否包含语音段时间戳
                if "timestamp" in result and result["timestamp"]:
                    # 检测到语音
                    if not self.is_speaking:
                        self.is_speaking = True
                        self.silence_counter = 0
                        logger.debug("VAD detected speech start")
                        return {"开始": True}
                    self.silence_counter = 0
                else:
                    # 未检测到语音
                    if self.is_speaking:
                        self.silence_counter += self.frame_time_ms
                        # 如果静默时间超过阈值，标记语音结束
                        if self.silence_counter >= self.min_silence_duration_ms:
                            self.is_speaking = False
                            self.silence_counter = 0
                            logger.debug("VAD detected speech end")
                            return {"end": True}

            # 没有特殊状态变化，返回 None
            return None

        except Exception as e:
            logger.error(f"Error in FSMN VAD processing: {e}")
            return None

    def reset_states(self):
        """
        重置 VAD 状态
        """
        try:
            # 重置模型缓存和内部状态
            self.cache = {}
            self.is_speaking = False
            self.silence_counter = 0
            logger.debug("FSMN VAD states reset.")
        except Exception as e:
            logger.error(f"Error resetting FSMN VAD states: {e}")


def create_instance(class_name, *args, **kwargs):
    """
    创建 VAD 实例的工厂函数
    Args:
        class_name: 类名
        *args: 位置参数
        **kwargs: 关键字参数
    Returns:
        实例化的对象
    """
    # 获取类对象
    cls = globals().get(class_name)
    if cls:
        # 创建并返回实例
        return cls(*args, **kwargs)
    else:
        raise ValueError(f"Class {class_name} not found")


# 测试代码
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.DEBUG)

    # 创建默认配置
    vad_config = {
        "sampling_rate": 16000,
        "threshold": 0.5,
        "min_silence_duration_ms": 200
    }

    # 创建 FSMNVAD 实例
    vad = FSMNVAD(vad_config)
    print("FSMN VAD instance created successfully")

    # 模拟一些音频数据进行测试
    try:
        # 生成随机音频数据（模拟语音）
        print("Generating test audio data...")
        speech_data = np.random.randint(-32768, 32767, size=16000, dtype=np.int16).tobytes()

        # 测试语音检测
        print("Testing VAD with speech data...")
        result = vad.is_vad(speech_data)
        print(f"VAD result: {result}")

        # 生成静默数据
        print("Generating silence data...")
        silence_data = np.zeros(16000, dtype=np.int16).tobytes()

        # 测试静默检测
        print("Testing VAD with silence data...")
        result = vad.is_vad(silence_data)
        print(f"VAD result: {result}")

        # 重置状态
        print("Resetting VAD states...")
        vad.reset_states()

    except Exception as e:
        print(f"Error during test: {e}")